import requests
import functools
import asyncio
from concurrent.futures import ThreadPoolExecutor

from datetime import datetime, timedelta
from core import setting, mgClient
import pymongo as mg
db = mgClient['chat']
draw_collection = db.get_collection('draw')


def check_traffic(ip):
    """
    It checks the traffic of the given IP address.
    
    :param ip: The IP address of the host to check
    """
    now = datetime.now()
    if ip != "221.215.48.6":
        traffic = draw_collection.find_one({"ip": ip})  # 查找指定IP地址的流量信息
        if not traffic:
            traffic = {"ip": ip, "data": [{"timestamp": now, "count": 1}]}
            draw_collection.insert_one(traffic)  # 插入新的流量信息
            return False
        else:
            recent_count = sum([t["count"] for t in traffic["data"]])
            # 删除超过24小时的旧流量数据
            traffic["data"] = [t for t in traffic["data"] if now - t["timestamp"] <= timedelta(hours=24)]
            if recent_count >= 100:  # 设置流量限制为10分钟内最多60个请求
                return True
            else:
                traffic["data"].append({"timestamp": now, "count": 1})
                draw_collection.update_one({"ip": ip}, {"$set": traffic})  # 更新流量信息
                return False
    else:
        return False
def get_api_key(api_keys, api_key_index):
    """
    It returns the next API key in the list of API keys
    :return: The api_key is being returned.
    """
    print(api_key_index)
    api_key = api_keys[api_key_index]
    api_key_index = (api_key_index + 1) % len(api_keys)
    return api_key, api_key_index

def get_api_keys(key_file_path):
    """
    It opens the file at the path specified by the argument key_file_path, reads each line of the
    file, strips the newline character from each line, and returns a list of the lines
    
    :param key_file_path: The path to the file containing the API keys
    :return: A list of API keys
    """
    api_keys = []
    with open(key_file_path, 'r') as f:
        for line in f:
            api_keys.append(line.strip())
    return api_keys

async def async_request(prompt: str, api_key: str):
    
    data = {
        "model": "gpt-3.5-turbo-16k",
        "messages": prompt,
        "temperature": 1,
        # "max_tokens": 4096,
        "stream": True,  # Enable streaming API
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": "Bearer "+api_key,
    }
    url = "https://api.openai.com/v1/chat/completions"
    loop = asyncio.get_running_loop()
    request = functools.partial(requests.post, url, headers=headers, json=data, stream = True, timeout = 3)
    try:
        with ThreadPoolExecutor() as executor:
            response = await loop.run_in_executor(
                executor, request
            )
    except:
        response = requests.Response()
        response.status_code = 500
    return response
